package tc.alert.tckejisoc.es;

import com.google.common.collect.Lists;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.IndexQuery;
import org.springframework.data.elasticsearch.core.query.IndexQueryBuilder;
import org.springframework.stereotype.Service;
import tc.alert.tckejisoc.domain.EventBase;

import java.util.List;
import java.util.stream.Collectors;


/**
 * 批量模式
 */

@Service
public class EventBulkWriter {
    @Autowired
    private ElasticsearchRestTemplate restTemplate;

    private static final int BATCH_SIZE = 2000; // 每批处理量

    public void bulkInsert(List<EventBase> events, String index) {
        List<List<EventBase>> batches = Lists.partition(events, BATCH_SIZE);

        batches.forEach(batch -> {
            List<IndexQuery> queries = batch.stream()
                    .map(event -> new IndexQueryBuilder()
                            .withId(event.getEventId())
                            .withObject(event)
                            .build())
                    .collect(Collectors.toList());

            restTemplate.bulkIndex(queries, IndexCoordinates.of(index));
        });
    }
}